import torch
import numpy as np
import csv

###gir_mirb_r50_triloss_ins.20240215
instance_file = "/root/ws/logs/models/girformer_gir_meshb_r50_160k_triloss/instance_features_final.pth"
out_csv = "/root/ws/logs/models/girformer_gir_meshb_r50_160k_triloss/analyze.csv"
###gir multiloss
# instance_file = "/root/ws/logs/tf-logs/gir_r50_ins/instance_features.pth"
# instance_file = "/root/ws/logs/tf-logs/gir_r50_ins/instance_features.pth"
# instance_file = "/root/autodl-tmp/models/girformer_gir_meshb_r50_160k_triloss/instance_features.pth"
def get_max_dis(arr):
    
    min_dis = 9999999999999.0
    length = len(arr)
    for ind in range(length):
        v = arr[ind]
        for ind2 in range(length):
            v2 = arr[ind2]
            if ind == ind2:
                continue
            dis = torch.linalg.norm(v-v2)    
            if dis < min_dis:
                min_dis = dis

    return min_dis

#获取每个类内部，距离中心点最远的点
def get_maxdis2center_inclass(res,ids):
    max_dis = 0
    all_features = res["all_features"]
    all_features_inds = res["all_ids"]
    centers = res["mean"]
    
    result = {}
    
    for id in ids:
        id_center = centers[id]
        id_inds = torch.where(all_features_inds == id)
        id_features = all_features[id_inds]
        id_fs_l2 = torch.linalg.norm(id_features-id_center,dim=1)
        
        dis = torch.max(id_fs_l2)
        result[id] = dis.item()
    return result

#获取每个类外部，距离该类中心点最远的点
def get_mindis2center_outclass(res,ids):
    max_dis = 0
    all_features = res["all_features"]
    all_features_inds = res["all_ids"]
    centers = res["mean"]
    
    result = {}
    
    for id in ids:
        id_center = centers[id]
        id_inds = torch.where(all_features_inds != id)
        id_features = all_features[id_inds]
        id_fs_l2 = torch.linalg.norm(id_features-id_center,dim=1)
        
        dis = torch.min(id_fs_l2)
        result[id] = dis.item()
    return result

#获取每个类中心点,距离其它类中心最近点的距离
def get_mindis_center2center(res,ids):
    min_dis_map = {}
    g_min_dis = 9999999999999.0
    centers = res["mean"]
    for id,v in centers.items():
        min_dis = 9999999999999.0
        for did in ids:
            if id == did:
                continue
            dv = centers[did]
            dis = torch.linalg.norm(v-dv)    
            if dis < min_dis:
                min_dis = dis
            if min_dis < g_min_dis:
                g_min_dis = min_dis
        min_dis_map[id] = min_dis.item()
    return min_dis_map

def do_eval(file):
    res = torch.load(file)
    ids = torch.unique(res["all_ids"]).numpy()
    
    min_out_dis = get_mindis2center_outclass(res,ids)
    max_in_dis = get_maxdis2center_inclass(res,ids)
    min_c2c_dis = get_mindis_center2center(res,ids)
    
    alldata = []
    alldata.append(["id","max_in_dis","min_out_dis","min_c2c_dis"])
    for id in ids:
        data = []
        data.append(id)
        data.append(max_in_dis[id])
        data.append(min_out_dis[id])
        data.append(min_c2c_dis[id])
        alldata.append(data)
    
    with open(out_csv, "w") as f:
        writer = csv.writer(f)
        writer.writerows(alldata)

do_eval(instance_file)


def do_analyze_file(file):
            # res = {
            # "mean":result_mean,
            # "std":result_std,
            # "all_features":g_all_feats_tensor,
            # "all_ids":g_all_ids_tensor
            # }
    res = torch.load(file)
    ids = torch.unique(res["all_ids"]).numpy()
       
    ###获取其他类中心到本类中心的最小距离
    min_dis_map = {}
    g_min_dis = 9999999999999.0
    centers = res["mean"]
    for id,v in centers.items():
        min_dis = 9999999999999.0
        for did in ids:
            if id == did:
                continue
            dv = centers[did]
            dis = torch.linalg.norm(v-dv)    
            if dis < min_dis:
                min_dis = dis
            if min_dis < g_min_dis:
                g_min_dis = min_dis
        min_dis_map[id] = min_dis
    print(f"MIN OUT CLASS CENTER DIS : {g_min_dis}")

    ###get max std
    stds = res["std"]
    stds = [v.item() for id,v in stds.items()]
    max_std = np.max(stds)
    print(f"MAX STD : {max_std}")
    
    ###获取每个ID(类)内部中所有特征间的最大距离
    max_dis = 0
    all_features = res["all_features"]
    all_features_inds = res["all_ids"]
    for id in ids:
        id_inds = torch.where(all_features_inds == id)
        id_features = all_features[id_inds]
        dis = get_max_dis(id_features)
        if dis > max_dis:
            max_dis = dis
    print(f"获取每个ID(类)内部中所有特征间的最大距离 MAX IN CLASS DIS : {max_dis}")

    ###获取不同类样本之间的最小距离
    min_dis_map = {}
    g_min_dis = 9999999999999.0
    centers = res["mean"]
    for id,v in centers.items():
        min_dis = 9999999999999.0
        for did in ids:
            if id == did:
                continue
            dv = centers[did]
            dis = torch.linalg.norm(v-dv)    
            if dis < min_dis:
                min_dis = dis
            if min_dis < g_min_dis:
                g_min_dis = min_dis
        min_dis_map[id] = min_dis

    


# do_analyze_file(instance_file)